# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import abc
import logging
from collections import namedtuple

from coin.strategy.mm.subscription import SubscriptionRequest
from coin.proto.coin_query_pb2 import (
    CurrencyBalance,
    AccountBalance,
)
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.exchange.binance_futures.kr_rest.futures_product import BinanceFuturesProduct

ReserveSubStratId = namedtuple('ReserveSubStratId', ['reserve', 'sub_strat_id'])


def create_pnl_balance_calculator(sub_req: SubscriptionRequest, quote):
  if sub_req.market_type == 'Spot':
    return SpotPnlBalanceCalculator(sub_req, quote)
  elif sub_req.market_type == 'Futures' and sub_req.exchange == 'Bitflyer':
    # Only support spot in Bitflyer.
    return SpotPnlBalanceCalculator(sub_req, quote)
  elif (sub_req.market_type == 'Futures'
        and sub_req.exchange in ['Binance', 'Okex', 'Bitmex', 'Huobi', 'Kraken', 'Deribit']):
    return FuturesPnlBalanceCalculator(sub_req, quote)
  else:
    raise NotImplementedError(sub_req)


class PnlBalanceCalculator(metaclass=abc.ABCMeta):
  def __init__(self, sub_req: SubscriptionRequest, quote):
    self._sub_req = sub_req
    self._quote = quote
    self._og = None
    self._product_map = {}
    self._books = {}
    self._logger = logging.getLogger('PnlBalanceCalculator')

  @abc.abstractmethod
  def calc_pnl_balance(self):
    raise NotImplementedError()

  def is_book_ready(self):
    # books contains price infos of all products in product_map
    return set(self._product_map.keys()) <= set(self._books.keys())

  def add_product_reserve_sub_strat_id_map(self, product, reserve, sub_strat_id=None):
    self._product_map[product] = ReserveSubStratId(reserve, sub_strat_id)

  def add_og(self, og):
    self._og = og

  def on_book(self, book):
    self._books[book.product] = book


class SpotPnlBalanceCalculator(PnlBalanceCalculator):
  def __init__(self, sub_req: SubscriptionRequest, quote):
    super().__init__(sub_req, quote)
    self._quote_proto = CurrencyBalance(
        currency=self._quote,
        currency_native=self._quote,
        quote_price=1,
    )

  def _calc_position_estimate(self, product):
    book = self._books[product]
    mid_price = (book.bid0().price + book.ask0().price) / 2
    balance = self._og.get_balance(product.base)
    pe = (balance - self._product_map[product].reserve) * mid_price
    return pe

  def calc_pnl_balance(self):
    if not self.is_book_ready():
      self._logger.error('Book data is not ready for PnL calculation!')
      return []

    pe_sum = 0
    for product in self._product_map.keys():
      pe_sum += self._calc_position_estimate(product)
    pnl_balance = self._og.get_balance(self._quote) + pe_sum

    pnl_balance_proto = CurrencyBalance(
        currency=self._quote,
        currency_native=self._quote,
        total=pnl_balance,
    )
    pnl_balance_result = StrategyLog(type=StrategyLog.PNL_BALANCE,
                                     pnl_balance=pnl_balance_proto,
                                     quote=self._quote_proto,
                                     subreq_id=None)
    return [pnl_balance_result]


class FuturesPnlBalanceCalculator(PnlBalanceCalculator):
  def __init__(self, sub_req: SubscriptionRequest, quote):
    super().__init__(sub_req, quote)
    self._og_info = None

  def _get_pnl_quote_price(self, product):
    book = self._books[product]
    mid_price = (book.bid0().price + book.ask0().price) / 2
    return mid_price

  # Inverse futures
  def calc_pnl_balance(self):
    if not self.is_book_ready():
      self._logger.error('Book data is not ready for PnL calculation!')
      return []

    pnl_balance_results = []
    for product in self._product_map.keys():
      if isinstance(product, BinanceFuturesProduct):
        pnl_quote = product.quote
      else:
        pnl_quote = product.base
      pnl_balance = self._og_info.get_margin_balance(pnl_quote.currency)
      pnl_balance_proto = CurrencyBalance(
          currency=pnl_quote.currency,
          currency_native=pnl_quote.native_currency,
          total=pnl_balance,
      )
      quote_proto = CurrencyBalance(
          currency=pnl_quote.currency,
          currency_native=pnl_quote.native_currency,
          quote_price=self._get_pnl_quote_price(product),
      )
      pnl_balance_results.append(
          StrategyLog(type=StrategyLog.PNL_BALANCE,
                      pnl_balance=pnl_balance_proto,
                      quote=quote_proto,
                      subreq_id=self._product_map[product].sub_strat_id))
    return pnl_balance_results

  def _find_og_info(self):
    if self._sub_req.market_type == 'Futures':
      if self._sub_req.exchange == 'Okex':
        self._og_info = self._og._og_info_rest
      else:
        self._og_info = self._og._og_info
    else:
      raise ValueError(self._sub_req)

  def add_og(self, og):
    super().add_og(og)
    self._find_og_info()
